package com.lingjtx.common.thread.commom;

import com.lingjtx.common.core.RedisOptions;
import com.lingjtx.common.core.ResponseResult;
import com.lingjtx.common.core.constant.SignatureType;
import com.lingjtx.common.core.util.SignatureHelper;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.server.WebFilter;
import org.springframework.web.server.WebFilterChain;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.nio.charset.StandardCharsets;
import java.util.Optional;

@Component
public class ApiSignatureFilter implements WebFilter {

    private final RedisOptions redis;

    public ApiSignatureFilter(RedisOptions redis) {
        this.redis = redis;
    }

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {
        return DataBufferUtils.join(exchange.getRequest().getBody())
                .flatMap(dataBuffer -> {
                    byte[] bodyBytes = new byte[dataBuffer.readableByteCount()];
                    dataBuffer.read(bodyBytes);
                    DataBufferUtils.release(dataBuffer);

                    // 重新构造 request 供后续读取
                    ServerHttpRequest decoratedRequest = new ServerHttpRequestDecorator(exchange.getRequest()) {
                        @Override
                        public Flux<DataBuffer> getBody() {
                            return Flux.just(exchange.getResponse().bufferFactory().wrap(bodyBytes));
                        }
                    };

                    ServerWebExchange mutatedExchange = exchange.mutate().request(decoratedRequest).build();

                    String signature = Optional.ofNullable(exchange.getRequest().getHeaders().getFirst("Signature")).orElse("");
                    String body = new String(bodyBytes, StandardCharsets.UTF_8);
                    SignatureType result = SignatureHelper.verify(signature, body, redis);
                    if (!SignatureType.SUCCESS.equals(result)) {
                        return writeErrorResponse(exchange, ResponseResult.createByError(401, result.toString()));
                    }

                    return chain.filter(mutatedExchange);
                });
    }

    private Mono<Void> writeErrorResponse(ServerWebExchange exchange, ResponseResult<String> error) {
        ServerHttpResponse response = exchange.getResponse();
        DataBuffer buffer = response.bufferFactory().wrap(error.toString().getBytes(StandardCharsets.UTF_8));
        return response.writeWith(Mono.just(buffer));
    }
}
